import os  # 操作系统接口
import pandas as pd  # 数据处理和分析
import glob  # 文件路径模式匹配
import argparse  # 命令行参数解析
from typing import List  # 类型提示

def get_available_classes(data_dir: str) -> List[str]:
    """
    Get the list of available classes from the data directory.
    
    Args:
        data_dir: Path to directory containing data files
        
    Returns:
        List of available class names
    """
    # Get all CSV files in the directory
    csv_files = glob.glob(os.path.join(data_dir, "*.csv"))  # 查找所有CSV文件
    
    if not csv_files:  # 检查是否找到文件
        raise FileNotFoundError(f"No CSV files found in {data_dir}")  # 抛出异常
    
    # Extract class names from filenames
    class_names = set()  # 使用集合避免重复
    for file_path in csv_files:  # 遍历所有CSV文件
        filename = os.path.basename(file_path)  # 获取文件名
        parts = filename.split('-')  # 按'-'分割文件名
        
        if len(parts) >= 3:  # 确保文件名格式正确
            class_name = parts[2]  # 提取类名(文件名第三部分)
            class_names.add(class_name)  # 添加到集合
    
    return list(class_names)

def prepare_data_files(data_dir: str, output_dir: str):
    """
    Prepare data files by copying them to the output directory.
    
    Args:
        data_dir: Path to directory containing original data files
        output_dir: Path to directory to save prepared data files
    """
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)  # 创建输出目录(如果不存在)
    
    # Get all CSV files in the directory
    csv_files = glob.glob(os.path.join(data_dir, "*.csv"))  # 查找所有CSV文件
    
    if not csv_files:  # 检查是否找到文件
        raise FileNotFoundError(f"No CSV files found in {data_dir}")  # 抛出异常
    
    # Process each CSV file
    for file_path in csv_files:  # 遍历每个CSV文件
        # Read the CSV file
        df = pd.read_csv(file_path)  # 读取CSV文件
        
        # Convert timestamp to datetime if 'ts' column exists
        if 'ts' in df.columns:  # 检查是否有时间戳列
            df['timestamp'] = pd.to_datetime(df['ts'], unit='ms')  # 转换为datetime
            df = df.sort_values('timestamp')  # 按时间排序
        
        # Get the output filename
        output_file = os.path.join(output_dir, os.path.basename(file_path))  # 构建输出路径
        
        # Save to CSV
        df.to_csv(output_file, index=False)  # 保存处理后的数据
        print(f"Saved {len(df)} samples to {output_file}")  # 打印保存信息

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Prepare data files for training')
    parser.add_argument('--data_dir', type=str, default='../data',
                       help='Path to directory containing original data files')
    parser.add_argument('--output_dir', type=str, default='../data_prepared',
                       help='Path to directory to save prepared data files')
    
    args = parser.parse_args()
    
    # Get available classes
    classes = get_available_classes(args.data_dir)
    print(f"Found {len(classes)} classes: {classes}")
    
    # Prepare data files
    prepare_data_files(args.data_dir, args.output_dir)

